import sys
import os

import torch
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QTabWidget, QLabel, QLineEdit, QPushButton, QTextEdit,
                             QFileDialog, QTableWidget, QTableWidgetItem, QScrollArea,
                             QComboBox, QMessageBox, QHeaderView)
from PyQt5.QtCore import Qt, QMimeData, pyqtSignal
from PyQt5.QtGui import QPixmap, QDragEnterEvent, QDropEvent
import pandas as pd
import numpy as np
from datasets import tqdm
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

from Config import Paths, Dconfig, device
from Data import Dataset
from Model import Model
from Utils import sigmoid, LOGGER
import warnings

warnings.filterwarnings("ignore", category=UserWarning)

def combined_inference(config, test_df, tokenizer, single_model, multi_fold_models, is_ensemble=False, is_fold=False, n_folds=5):
    test_dataset = Dataset(config, test_df, tokenizer)

    test_loader = DataLoader(
        test_dataset,
        batch_size=config.BATCH_SIZE_TEST,
        shuffle=False,
        num_workers=config.NUM_WORKERS,
        pin_memory=True,
        drop_last=False
    )

    all_ids = []
    with torch.no_grad():
        for batch in test_loader:
            all_ids.extend(batch["ids"])
    all_ids = np.array(all_ids)

    if is_ensemble:
        all_preds = []
        for model in multi_fold_models:
            fold_preds = []
            with torch.no_grad(), tqdm(test_loader, desc="Inference") as pbar:
                for batch in pbar:
                    inputs = {k: v.to(device) for k, v in batch["inputs"].items()}
                    preds = model(inputs).cpu().numpy().flatten()
                    fold_preds.append(preds)
            all_preds.append(np.concatenate(fold_preds))
        assert len(all_ids) == len(all_preds[0]), "IDs和预测值数量不匹配！"

        predictions = np.mean(all_preds, axis=0)
    else:
        all_preds = []
        with torch.no_grad(), tqdm(test_loader, desc="Inference") as pbar:
            for batch in pbar:
                inputs = {k: v.to(device) for k, v in batch["inputs"].items()}
                preds = single_model(inputs).cpu().numpy().flatten()
                all_preds.append(preds)

        predictions = np.concatenate(all_preds)

    return {
        "predictions": predictions,
        "ids": all_ids
    }

class DragDropWidget(QWidget):
    filesDropped = pyqtSignal(list)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setAcceptDrops(True)
        self.initUI()

    def initUI(self):
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)

        self.label = QLabel("拖拽文件到此处或点击选择文件")
        self.label.setAlignment(Qt.AlignCenter)
        self.label.setStyleSheet("""
            QLabel {
                border: 2px dashed #aaa;
                padding: 20px;
                font-size: 14px;
            }
        """)
        self.layout.addWidget(self.label)

        self.btn = QPushButton("选择文件")
        self.btn.clicked.connect(self.selectFiles)
        button_width = 160
        button_height = 90
        self.btn.setFixedSize(button_width, button_height)
        self.layout.addWidget(self.btn, alignment=Qt.AlignCenter)

        font = self.btn.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.btn.setFont(font)

    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QDropEvent):
        files = [url.toLocalFile() for url in event.mimeData().urls()]
        self.filesDropped.emit(files)
        event.acceptProposedAction()

    def selectFiles(self):
        files, _ = QFileDialog.getOpenFileNames(self, "选择文件", "",
                                                "CSV Files (*.csv);;Text Files (*.txt);;All Files (*)")
        if files:
            self.filesDropped.emit(files)


class TextInputWidget(QWidget):
    textSubmitted = pyqtSignal(str)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.initUI()

    def initUI(self):
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)

        self.label = QLabel("输入待分析文本:")
        font = self.label.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.label.setFont(font)
        self.layout.addWidget(self.label)

        self.textEdit = QTextEdit()
        self.textEdit.setPlaceholderText("在此输入文本内容...")
        self.layout.addWidget(self.textEdit)

        self.btn = QPushButton("分析文本")
        self.btn.clicked.connect(self.analyzeText)
        button_width = 160
        button_height = 90
        self.btn.setFixedSize(button_width, button_height)
        self.layout.addWidget(self.btn, alignment=Qt.AlignCenter)

        font = self.btn.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.btn.setFont(font)

    def analyzeText(self):
        text = self.textEdit.toPlainText().strip()
        if text:
            self.textSubmitted.emit(text)
        else:
            QMessageBox.warning(self, "警告", "请输入文本内容!")


class ResultTableWidget(QTableWidget):
    """结果显示表格组件"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setColumnCount(2)  # 修改为两列，只显示ID和分类结果
        self.setHorizontalHeaderLabels(["ID", "分类结果"])
        self.horizontalHeader().setStretchLastSection(True)
        self.setEditTriggers(QTableWidget.NoEditTriggers)
        self.setSelectionBehavior(QTableWidget.SelectRows)

        header = self.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.Stretch)
        header.setSectionResizeMode(1, QHeaderView.Stretch)

    def displayResults(self, results):
        """显示结果数据"""
        self.clearContents()
        self.setRowCount(len(results))

        for row, (id_, pred) in enumerate(results):
            self.setItem(row, 0, QTableWidgetItem(str(id_)))
            classification = "AI生成" if pred > Paths.BEST_Ensemble_THRESHOLD else "人类写作"
            self.setItem(row, 1, QTableWidgetItem(classification))

        header = self.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.Stretch)
        header.setSectionResizeMode(1, QHeaderView.Stretch)


class VisualizationWidget(QWidget):
    """可视化组件"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.initUI()

    def initUI(self):
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)

        self.label = QLabel("预测结果分布")

        font = self.label.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.label.setFont(font)
        self.layout.addWidget(self.label)

        width = 8
        height = width * 9 / 16
        self.figure = plt.figure(figsize=(width, height))
        self.canvas = FigureCanvas(self.figure)
        self.layout.addWidget(self.canvas)

    def updateVisualization(self, predictions, is_ensemble):
        plt.rcParams['font.sans-serif'] = ['SimHei']

        plt.rcParams['axes.unicode_minus'] = False


        """更新可视化图表"""
        self.figure.clear()
        ax = self.figure.add_subplot(111)

        if is_ensemble:
            ai_count = sum([1 for pred in predictions if pred > Paths.BEST_Ensemble_THRESHOLD])
            human_count = len(predictions) - ai_count
            labels = ['AI生成占比', '人类写作占比']
            sizes = [ai_count, human_count]
            ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=90)
        else:
            ai_prob = predictions[0]
            human_prob = 1 - ai_prob
            labels = ['AI生成可能性', '人类写作可能性']
            sizes = [ai_prob, human_prob]
            ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=90)

        ax.axis('equal')
        self.canvas.draw()


class FileTab(QWidget):
    """文件输入标签页，包含状态提示"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.initUI()

    def initUI(self):
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)

        self.statusLabel = QLabel("准备就绪")

        font = self.statusLabel.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.statusLabel.setFont(font)
        self.statusLabel.setAlignment(Qt.AlignCenter)
        self.statusLabel.setStyleSheet("""
            QLabel {
                font-weight: bold;
                color: #666;
                padding: 5px;
                border-bottom: 1px solid #eee;
            }
        """)
        self.layout.addWidget(self.statusLabel)

        self.dragDropWidget = DragDropWidget()
        self.layout.addWidget(self.dragDropWidget)

        self.resultTable = ResultTableWidget()
        self.layout.addWidget(self.resultTable)


class MainWindow(QMainWindow):
    """主窗口"""

    def __init__(self):
        super().__init__()
        self.setWindowTitle("文本分类检测工具")
        self.resize(1000, 700)
        self.initUI()

        self.tokenizer = AutoTokenizer.from_pretrained(os.path.join(Paths.OUTPUT_DIR, "tokenizer"))

        self.single_model = Model(Dconfig, config_path=Paths.CONFIG_PATH, pretrained=False, is_training=False)
        self.single_model.load_state_dict(torch.load(f"./output/fold_0_best.pth"))
        self.single_model.to(device)
        self.single_model.eval()

        self.multi_fold_models = []
        n_folds = 5
        for fold in range(n_folds):
            model_path = os.path.join(Paths.OUTPUT_DIR, f"fold_{fold}_best.pth")

            if not os.path.exists(model_path):
                LOGGER.warning(f"Model not found: {model_path}")
                continue

            model = Model(Dconfig, config_path=Paths.CONFIG_PATH, pretrained=False, is_training=False)
            model.load_state_dict(torch.load(model_path))
            model.to(device)
            model.eval()
            self.multi_fold_models.append(model)

    def initUI(self):

        mainWidget = QWidget()
        mainLayout = QVBoxLayout()
        mainWidget.setLayout(mainLayout)
        self.setCentralWidget(mainWidget)

        modeLayout = QHBoxLayout()
        self.modeLabel = QLabel("检测模式:")

        font = self.modeLabel.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.modeLabel.setFont(font)
        self.modeCombo = QComboBox()
        self.modeCombo.addItems(["快速检测", "综合检测"])
        modeLayout.addWidget(self.modeLabel)
        modeLayout.addWidget(self.modeCombo)
        modeLayout.addStretch()
        mainLayout.addLayout(modeLayout)

        self.tabWidget = QTabWidget()
        mainLayout.addWidget(self.tabWidget)

        self.textTab = QWidget()
        self.textTabLayout = QVBoxLayout()
        self.textTab.setLayout(self.textTabLayout)

        self.textInput = TextInputWidget()
        self.textInput.textSubmitted.connect(self.processSingleText)
        self.textTabLayout.addWidget(self.textInput)

        self.singleResultLabel = QLabel("预测结果:")

        font = self.singleResultLabel.font()
        font.setPointSize(int(font.pointSize() * 1.5))
        self.singleResultLabel.setFont(font)
        self.textTabLayout.addWidget(self.singleResultLabel)

        self.singleResultDisplay = QLabel("请输入文本")
        self.textTabLayout.addWidget(self.singleResultDisplay)

        self.tabWidget.addTab(self.textTab, "文本输入")

        self.fileTab = FileTab(self)  # 使用新的FileTab类
        self.fileTab.dragDropWidget.filesDropped.connect(self.processFiles)
        self.tabWidget.addTab(self.fileTab, "文件输入")

        self.visualTab = QWidget()
        self.visualTabLayout = QVBoxLayout()
        self.visualTab.setLayout(self.visualTabLayout)

        self.visualization = VisualizationWidget()
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setWidget(self.visualization)
        self.visualTabLayout.addWidget(scroll)

        self.tabWidget.addTab(self.visualTab, "结果可视化")

    def processSingleText(self, text):
        """处理单个文本输入"""
        self.singleResultDisplay.setText("检测中...")
        QApplication.processEvents()

        temp_id = str(np.random.randint(100000))
        df = pd.DataFrame({
            "id": [temp_id],
            "text": [text],
            "generated": [0]
        })

        ensemble = self.modeCombo.currentText() == "复合模型"

        try:
            predictions = combined_inference(
                Dconfig, df, self.tokenizer,
                self.single_model, self.multi_fold_models,
                is_ensemble=ensemble,
                n_folds=5
            )

            pred_value = sigmoid(predictions["predictions"][0])
            classification = "AI生成" if pred_value > 0.5 else "人类写作"
            self.singleResultDisplay.setText(
                f"{classification}"
            )

            self.visualization.updateVisualization([pred_value], is_ensemble=False)

        except Exception as e:
            QMessageBox.critical(self, "错误", f"检测过程中发生错误:\n{str(e)}")
            self.singleResultDisplay.setText("分析失败，请重试")

    def processFiles(self, file_paths):
        """处理文件输入"""

        self.fileTab.statusLabel.setText("检测中...")
        self.fileTab.statusLabel.setStyleSheet("font-weight: bold; color: #FF9900;")
        QApplication.processEvents()

        dfs = []
        for file_path in file_paths:
            try:
                if file_path.endswith('.csv'):
                    df = pd.read_csv(file_path)
                else:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        texts = f.readlines()
                    df = pd.DataFrame({
                        "id": [str(i) for i in range(len(texts))],
                        "text": texts,
                        "generated": [0] * len(texts)
                    })

                if "text" not in df.columns:
                    df["text"] = df.iloc[:, 0]

                if "id" not in df.columns:
                    df["id"] = [str(i) for i in range(len(df))]

                dfs.append(df)

            except Exception as e:
                QMessageBox.warning(self, "警告", f"无法读取文件 {file_path}:\n{str(e)}")
                continue

        if not dfs:
            self.fileTab.statusLabel.setText("没有有效的文件可处理!")
            self.fileTab.statusLabel.setStyleSheet("font-weight: bold; color: #FF0000;")
            self.fileTab.resultTable.setRowCount(0)
            return

        combined_df = pd.concat(dfs, ignore_index=True)
        ensemble = self.modeCombo.currentText() == "复合模型"

        try:
            predictions = combined_inference(
                Dconfig, combined_df, self.tokenizer,
                self.single_model, self.multi_fold_models,
                is_ensemble=ensemble,
                n_folds=5
            )

            results = list(zip(predictions["ids"], sigmoid(predictions["predictions"])))
            self.fileTab.resultTable.displayResults(results)
            self.visualization.updateVisualization(sigmoid(predictions["predictions"]), is_ensemble=True)

            self.fileTab.statusLabel.setText(f"分析完成，共处理 {len(results)} 条数据")
            self.fileTab.statusLabel.setStyleSheet("font-weight: bold; color: #009900;")

            self.tabWidget.setCurrentIndex(1)

        except Exception as e:
            self.fileTab.statusLabel.setText("分析失败，请重试")
            self.fileTab.statusLabel.setStyleSheet("font-weight: bold; color: #FF0000;")
            self.fileTab.resultTable.setRowCount(0)
            QMessageBox.critical(self, "错误", f"检测过程中发生错误:\n{str(e)}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())